'''
Created on May 24, 2011

@author: erik
'''

from tcdb.api import Get
from tcdb.utils import Utils
from Bio.Blast import NCBIWWW
from Bio.Blast import NCBIXML
from Bio import SeqIO
from Bio import Entrez
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC

class Creator(object):
    Entrez.email = 'eclarke@ucsd.edu'
    '''
    This class creates a sample population through a number of different methods, 
    including database sampling and shuffling
    '''
    
    def __init__(self):
        pass
    
    def fromSeqFile(self, sequences_infile, e_val_max = 0.005, e_val_min = 0):
        
        self.e_val_max = e_val_max
        self.e_val_min = e_val_min
        population = []
        for seq_record in SeqIO.parse(sequences_infile, 'fasta', IUPAC.protein):
            population.append(seq_record)
            gi_num = []
            blast_handle = NCBIWWW.qblast("blastp", "nr", Utils.getId(seq_record), format_type="XML", expect=e_val_max)
            
            for blast_record in NCBIXML.parse(blast_handle):
                for alignment in blast_record.alignments:
                    for hsp in alignment.hsps:
                        gi_num.append(alignment.title.split("|")[1])
            entrez_handle = Entrez.efetch("protein", id=",".join(gi_num), rettype="fasta", retmode="xml")
            for entrez_record in Entrez.parse(entrez_handle):
                seq = Seq(entrez_record['TSeq_sequence'])
                new_seq_record = SeqRecord(seq)
                new_seq_record.id = entrez_record['TSeq_gi']
                population.append(new_seq_record)
        
        print population    
        return population
        
if __name__ == '__main__':
    create = Creator()
    create.fromSeqFile('/home/erik/workspace/tcdb/test.faa')
    